import numpy as np
import matplotlib.pyplot as plt
import pywt
from scipy.signal import stft, hilbert
from scipy.fftpack import fft

# 设置随机种子以确保每次运行代码生成相同的随机数
np.random.seed(0)

# 定义采样频率和时间向量
fs = 1000  # 采样频率
t = np.linspace(0, 1, fs, endpoint=False)

# 创建复杂信号，结合多个频率的正弦波和随机相位
amplitude = 1.5  # 基础振幅
signal = sum(amplitude * np.random.rand() * np.sin(2 * np.pi * np.random.randint(30, 100) * t + np.random.rand())
             for _ in range(5))  # 随机生成5个正弦波的和

# 添加白噪声
noise_amplitude = 0.5  # 噪声的振幅
noise = noise_amplitude * np.random.normal(size=t.shape)
mixed_signal = signal + noise  # 混合信号

# 第一阶段处理：小波变换去噪
coeffs = pywt.wavedec(mixed_signal, 'db4', level=4)
threshold = 0.1 * max(coeffs[0])  # 动态阈值设定为最大细节系数的10%
thresholded_coeffs = [pywt.threshold(i, threshold, mode='soft') for i in coeffs]
wavelet_signal = pywt.waverec(thresholded_coeffs, 'db4')

# 对去噪后的信号进行快速傅里叶变换 (FFT)
fft_signal = fft(wavelet_signal)
frequencies = np.fft.fftfreq(len(fft_signal), 1/fs)

# 第二阶段处理：SFFT和希尔伯特变换
f, t_, Zxx = stft(wavelet_signal, fs, nperseg=100)
hilbert_signal = hilbert(wavelet_signal)
envelope_signal = np.abs(hilbert_signal)

# 定义计算信噪比的函数
def calculate_snr(signal, noise):
    signal_power = np.mean(signal**2)
    noise_power = np.mean(noise**2)
    return 10 * np.log10(signal_power / noise_power)

# 计算信噪比
snr_wavelet = calculate_snr(wavelet_signal, mixed_signal - wavelet_signal)
snr_fft = calculate_snr(fft_signal, wavelet_signal - fft_signal)
snr_envelope = calculate_snr(envelope_signal, mixed_signal - envelope_signal)

# 信号功率
signal_power = np.mean(signal**2)
# 噪声功率
noise_power = np.mean(noise**2)


# 输出信号功率和噪声功率\
print(f"Signal Power: {signal_power:.2f}")
print(f"Noise Power: {noise_power:.2f}")
# 输出信噪比结果
print(f"SNR of original signal: {calculate_snr(signal, noise):.2f} dB")
print(f"SNR after wavelet processing: {snr_wavelet:.2f} dB")
print(f"SNR after FFT processing: {snr_fft:.2f} dB")
print(f"SNR after envelope extraction: {snr_envelope:.2f} dB")


